from antlr4 import *
from xscanner.grammar.RuleLexer import RuleLexer
from xscanner.grammar.RuleParser import RuleParser
from xscanner.modules import *


class Result:
    def __init__(self):
        self._name = ""  # 规则名
        self._rid = ""  # 规则ID = self.jls_extract_def()
        self._result = False  # 表示是否命中
        self._detail = ""  # 命中细节，哪个文件命中，哪个类命中。

    @property
    def name(self):
        return self._name

    @name.setter
    def name(self, value):
        self._name = value

    @property
    def rid(self):
        return self._rid

    @rid.setter
    def rid(self, value):
        self._rid = value

    @property
    def result(self):
        return self._result

    @result.setter
    def result(self, value):
        self._result = value

    @property
    def detail(self):
        return self._detail

    @detail.setter
    def detail(self, value):
        self._detail = value

    def __str__(self):
        return "{} {} {}".format(self._name, self._rid, self._result)


class VisitorImp(ParseTreeVisitor):
    def __init__(self, data: Data):
        self.data = data

        self._counter = {}  # 存放计数器 string:int
        self._hits = 0  # 存放命中的数量
        self._total = 0  # 存放规则总数
        # 存放临时变量，用于运算
        self._idents = []  # 存放变量名 _variables的key
        self._variables = {}  # 存放变量的值 string:bool

        self._is_skip = False  # 类型不匹配，则跳过。

        self._results = []  # 存放扫描结果
        self._result = None  # 存放一个扫描结果

    def get_result(self):
        return self._results

    def has(self, name):
        return name in self._variables

    def get(self, name):
        return self._variables.get(name, None)

    def visit(self, ctx: RuleParser.RulesContext):
        ctx.accept(self)

    # Visit a parse tree produced by RuleParser#rules.
    def visitRules(self, ctx: RuleParser.RulesContext):
        self.visitChildren(ctx)

    # Visit a parse tree produced by RuleParser#ruleDec.
    def visitRuleDec(self, ctx: RuleParser.RuleDecContext):
        self._result = Result()
        self._hits = 0
        self._total = 0

        rid = ctx.DECIMAL_LIT()
        s = []
        for item in ctx.IDENTIFIER():
            s.append(item.getText())
        name = ".".join(s)

        self._result.name = name
        self._result.rid = rid
        self.visitChildren(ctx)

    # Visit a parse tree produced by RuleParser#typeExpression.
    def visitTypeExpression(self, ctx: RuleParser.TypeExpressionContext):
        t = "any"
        if x := ctx.STRING_LIT():
            t = x.getText()[1:-1]

        if t == "any":
            return

        self._is_skip = t != self.data.type()

    def visitMainExpression(self, ctx: RuleParser.MainExpressionContext):
        if self._is_skip:
            return
        return self.visitChildren(ctx)

    # Visit a parse tree produced by RuleParser#lineExpression.
    def visitLineExpression(self, ctx: RuleParser.LineExpressionContext):
        if exp := ctx.assignExpression():
            result = self.visitAssignExpression(exp)
            if isinstance(result, bool):
                self._total += 1
                if result is True:
                    self._hits += 1  # 命中

            ident = exp.IDENTIFIER()
            if ident:
                self._idents.append(ident.getText())
                self._variables[ident.getText()] = result

        if exp := ctx.printExpression():
            return self.visitPrintExpression(exp)

    @classmethod
    def get_argument(cls, argumentExpression):
        if x := argumentExpression.STRING_LIT():
            return x.getText()[1:-1], MODE_STRING
        elif x := argumentExpression.STRING_REGEXP_LIT():
            return x.getText()[1:-1], MODE_STRING_REGEX
        elif x := argumentExpression.HEX_LIT():
            return x.getText()[1:-1], MODE_HEX
        elif x := argumentExpression.HEX_REGEXP_LIT():
            return x.getText()[1:-1], MODE_HEX_REGEX
        raise Exception("参数获取异常！")

    @classmethod
    def get_one(cls, oneArgumentExpression):
        arg0, flag0 = cls.get_argument(oneArgumentExpression.argumentExpression())
        return arg0, flag0

    @classmethod
    def get_two(cls, twoArgumentExpression):
        arg0, flag0 = cls.get_argument(twoArgumentExpression.argumentExpression(0))
        arg1, flag1 = cls.get_argument(twoArgumentExpression.argumentExpression(1))

        return arg0, flag0, arg1, flag1

    # Visit a parse tree produced by RuleParser#assignExpression.
    def visitAssignExpression(self, ctx: RuleParser.AssignExpressionContext):
        result = False

        if exp := ctx.zipSubFileExpression():
            arg0, flag0, arg1, flag1 = self.get_two(exp.twoArgumentExpression())
            return self.data.match_subfile(arg0, flag0, arg1, flag1)
        elif exp := ctx.apkManifestExpression():
            arg0, arg1 = self.get_one(exp.oneArgumentExpression())
            return self.data.match_manifest(arg0, arg1)
        elif exp := ctx.APK_MANIFEST_MAIN_ACTIVITY():
            return self.data.get_manifest_main_activities()
        elif exp := ctx.APK_MANIFEST_APPLICATION():
            return self.data.apk.get_manifest_application()
        elif exp := ctx.apkDexStringExpression():
            arg0, arg1 = self.get_one(exp.oneArgumentExpression())
            return self.data.match_dex_string(arg0, arg1)
        elif exp := ctx.apkDexClassExpression():
            return self.visitApkDexClassExpression(exp)
        elif exp := ctx.apkSubFileExpression():
            arg0, flag0, arg1, flag1 = self.get_two(exp.twoArgumentExpression())
            return self.data.match_subfile(arg0, flag0, arg1, flag1)
        else:
            arg0, arg1 = self.get_argument(ctx)
            return self.data.match(arg0, arg1)


    # Visit a parse tree produced by RuleParser#apkDexClassExpression.
    def visitApkDexClassExpression(self, ctx:RuleParser.ApkDexClassExpressionContext):
        if i := ctx.STRING_LIT():
            return self.data.match_dex_class(i.getText()[1:-1])
        
        if i := ctx.STRING_REGEXP_LIT():
            return self.data.match_dex_class(i.getText()[1:-1])
        
        if i := ctx.IDENTIFIER().getText():
            v = self._variables.get(i)
            if isinstance(v, list):
                for clz in v:
                    clz = clz.replace(".", "/")
                    if self.data.match_dex_class(clz):
                        return True
                        break
            else:
                return self.data.match_dex_class(v.replace(".", "/"))
        
        return False

    # Visit a parse tree produced by RuleParser#arithmeticExpression.
    def visitArithmeticExpression(self, ctx: RuleParser.ArithmeticExpressionContext):
        if i := ctx.inner:
            return self.visitArithmeticExpression(i)

        if i := ctx.arithmeticOperator():
            left = self.visitArithmeticExpression(ctx.left)
            right = self.visitArithmeticExpression(ctx.right)
            return left + right  # TODO 临时用加法

        if i := ctx.DECIMAL_LIT():
            return i.getText()

        if i := ctx.FLOAT_LIT():
            return i.getText()

        if i := ctx.IDENTIFIER():
            return self._variables.get(i.getText())

        raise Exception("未知语句")


    # Visit a parse tree produced by RuleParser#printExpression.
    def visitPrintExpression(self, ctx: RuleParser.PrintExpressionContext):
        print(self._variables.get(ctx.IDENTIFIER().getText()))


# ! ----------------------------------------------------------------

    # Visit a parse tree produced by RuleParser#ofExpression.
    def visitOfExpression(self, ctx: RuleParser.OfExpressionContext):
        hits = 0  # 命中数
        total = 0  # 总数

        if x := ctx.THEM():
            hits = self._hits
            total = self._total
        else:
            pass

        result = False
        if ctx.ANY():
            result = hits > 0
        elif ctx.ALL():
            result = hits == total
        else:
            # TODO 需要完善语法
            num = int(ctx.DECIMAL_LIT().getText())
            result = hits >= num

        return result

    # Visit a parse tree produced by RuleParser#conditionExpression.
    def visitConditionExpression(self, ctx: RuleParser.ConditionExpressionContext):
        if self.visitChildren(ctx):
            self._results.append(self._result)

    def visitLogicalExpression(self, ctx: RuleParser.LogicalExpressionContext):
        if i := ctx.inner:
            return self.visitLogicalExpression(i)
        if i := ctx.logicalOperator():
            left = self.visitLogicalExpression(ctx.left)
            right = self.visitLogicalExpression(ctx.right)
            if i.getText() == "&&":
                return left and right
            return left or right
        elif i := ctx.notExpression():
            return not self.visitLogicalExpression(i.logicalExpression())
        elif i := ctx.IDENTIFIER():
            return self._variables.get(i.getText())
        elif i := ctx.ofExpression():
            return self.visitOfExpression(i)
        elif i := ctx.relationExpression():
            return self.visitRelationExpression(i)

        raise Exception("未知语句")

    # Visit a parse tree produced by RuleParser#relationExpression.
    def visitRelationExpression(self, ctx: RuleParser.RelationExpressionContext):
        left = self.visitArithmeticExpression(ctx.left)
        right = self.visitArithmeticExpression(ctx.right)

        # 需要考虑 加减乘除的优先级，还要考虑括号问题
        return left == right  # TODO 临时用等于号
